/*
 * @Author: liusheng
 * @Date: 2022-06-04 18:43:11
 * @LastEditors: liusheng
 * @LastEditTime: 2022-06-04 23:21:01
 * @Description: 剑指 Offer II 071. 按权重生成随机数
 * email:liusheng613@126.com
 * Copyright (c) 2022 by liusheng/liusheng, All Rights Reserved. 
 * 
 * 剑指 Offer II 071. 按权重生成随机数
给定一个正整数数组 w ，其中 w[i] 代表下标 i 的权重（下标从 0 开始），请写一个函数 pickIndex ，它可以随机地获取下标 i，选取下标 i 的概率与 w[i] 成正比。

例如，对于 w = [1, 3]，挑选下标 0 的概率为 1 / (1 + 3) = 0.25 （即，25%），而选取下标 1 的概率为 3 / (1 + 3) = 0.75（即，75%）。

也就是说，选取下标 i 的概率为 w[i] / sum(w) 。

 

示例 1：

输入：
inputs = ["Solution","pickIndex"]
inputs = [[[1]],[]]
输出：
[null,0]
解释：
Solution solution = new Solution([1]);
solution.pickIndex(); // 返回 0，因为数组中只有一个元素，所以唯一的选择是返回下标 0。
示例 2：

输入：
inputs = ["Solution","pickIndex","pickIndex","pickIndex","pickIndex","pickIndex"]
inputs = [[[1,3]],[],[],[],[],[]]
输出：
[null,1,1,1,1,0]
解释：
Solution solution = new Solution([1, 3]);
solution.pickIndex(); // 返回 1，返回下标 1，返回该下标概率为 3/4 。
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 1
solution.pickIndex(); // 返回 0，返回下标 0，返回该下标概率为 1/4 。

由于这是一个随机问题，允许多个答案，因此下列输出都可以被认为是正确的:
[null,1,1,1,1,0]
[null,1,1,1,1,1]
[null,1,1,1,0,0]
[null,1,1,1,0,1]
[null,1,0,1,0,0]
......
诸若此类。
 

提示：

1 <= w.length <= 10000
1 <= w[i] <= 10^5
pickIndex 将被调用不超过 10000 次
 

注意：本题与主站 528 题相同： https://leetcode-cn.com/problems/random-pick-with-weight/

通过次数7,892   提交次数15,804
 */

#include "header.h"
class Solution {
public:
    Solution(vector<int>& w):generator(std::random_device{}()),distribution(1,std::accumulate(w.begin(),w.end(),0)),prefixSum(w.size(),0)
    {
        int total = 0; 
        for (int i = 0; i < w.size(); ++i)
        {
            total += w[i];
            prefixSum[i] = total;
        }
    }
    
    int pickIndex() {
        int val = distribution(generator);
        //lower bound 
        return std::lower_bound(prefixSum.begin(),prefixSum.end(),val) - prefixSum.begin();
    }
private:
    std::mt19937 generator;
    std::uniform_int_distribution<int> distribution;
    std::vector<int> prefixSum;
};

class Solution {
public:
    Solution(vector<int>& w):generator(std::random_device{}()),distribution(1,std::accumulate(w.begin(),w.end(),0)),prefixSum(w.size(),0)
    {
        int total = 0; 
        for (int i = 0; i < w.size(); ++i)
        {
            total += w[i];
            prefixSum[i] = total;
            // printf("prefixSum[%d]:%d\n",i,prefixSum[i]);
        }
    }
    
    int pickIndex() {
        int val = distribution(generator);
        
        //二分法在前缀和数组中找到第一个大于随机数的元素下标 lower_bound 
        int first = 0;
        int iter = 0;
        int step = 0;
        int count = prefixSum.size();
        while (count > 0)
        {
            iter = first;
            step = count / 2;
            iter += step;
            if (prefixSum[iter] < val)
            {
                first = ++iter;
                count -= step + 1; ///// === count -= (step + 1);
            }
            else
            {
                count = step;
            }
        }
        
        return first;
    }
private:
    std::mt19937 generator;
    std::uniform_int_distribution<int> distribution;
    std::vector<int> prefixSum;
};

